# 数据集验证，需要加载的包库
import os
from PIL import Image
from matplotlib import pyplot as plt
import xml
import cv2 as cv
import numpy as np
import json
import math
from scipy.spatial import distance
from collections import OrderedDict
import time
import pandas as pds
from torch.utils.data import DataLoader
from torchvision import models
import torch
import torch.nn as nn
from utilour.datasetLoader import MyDataset
import torch.optim as optimer
from logdir import loginfo
from utilour.tool import calaccurary,visualize_cam
from utilour.gradcadpp import GradCAM,GradCAMpp
import torchvision
from torchvision.utils import make_grid, save_image
import shutil
from torch.utils.tensorboard import SummaryWriter

def cleardir(foldirpath):
    if os.path.exists(foldirpath):
        shutil.rmtree(foldirpath)
    os.mkdir(foldirpath)


def loginfo(fpath,lslist,acclist,epoch):
    with open(fpath,'a',encoding="utf-8") as fp:
        for i in range(len(lslist)):
            fp.write("{},{},{}\n".format(epoch,lslist[i],acclist[i]))

'''
将训练和测试代码全部放到一个方法中，验证是模型的问题还是方法的问题
'''
def train(epoch,stepsum,model,dataloader,writer,modelcpk,traindecary,lr=0.01):
    model.train()
    lr=lr/((epoch//10)*10) if epoch>=10 else lr
    Optimizer=optimer.SGD(model.parameters(),lr=lr,momentum=0.9)
    CRLoss=nn.CrossEntropyLoss()

    lslist=[]
    aclist=[]
    Optimizer.zero_grad() # 导数归零
    for step,(data,label,imgpath) in  enumerate(dataloader):
        data,label=data.cuda(),label.cuda()
        tlabel=model(data) # 计算结果
        loss=CRLoss(tlabel,label)
        #if loss>traindecary:
        loss.backward()
        Optimizer.step()
        Optimizer.zero_grad() # 导数归零
        accurary=calaccurary(tlabel.detach(),label.detach())
        lslist.append(loss.detach().cpu().numpy())
        aclist.append(accurary[0])

        writer.add_scalar("train loss",loss,stepsum)
        writer.add_scalar("train acc",accurary[0],stepsum)
        stepsum=stepsum+1
    if epoch>30:    
        torch.save(model,os.path.join(modelcpk,"model_last.pkl"))
    return model,lslist,aclist,stepsum

def test(epoch,stepsum,model,dataloader,writer,lr=0.01,sig="vail"):
    model.eval()
    CRLoss=nn.CrossEntropyLoss()
    lslist=[]
    aclist=[]
    for step,(data,label,imgpath) in  enumerate(dataloader):
        
        data,label=data.cuda(),label.cuda()
        tlabel=model(data) # 计算结果
        loss=CRLoss(tlabel,label)
        accurary=calaccurary(tlabel.detach(),label.detach())
        lslist.append(loss.detach().cpu().numpy())
        aclist.append(accurary[0])
        writer.add_scalar("{} loss".format(sig),loss,stepsum)
        writer.add_scalar("{} acc".format(sig),accurary[0],stepsum)
        stepsum=stepsum+1
         
    return model,lslist,aclist,stepsum

def printinfo(epoch,ls,als,sig='train'):
    ls=np.array(ls)
    als=np.array(als)
    meanls,maxls,minls=np.mean(ls),np.max(ls),np.min(ls)
    meanas,maxas,minas=np.mean(als),np.max(als),np.min(als)
    print("epoch:{} meanacc:{} minacc:{} maxacc:{} meanloss:{} minloss:{} maxloss:{} sig:{}".format(epoch,
    "%.4f"%meanas,"%.4f"%minas,"%.4f"%maxas,
    "%.4f"%meanls,"%.4f"%minls,"%.4f"%maxls,
    sig
    ))
    return meanas,meanls


if __name__ == "__main__":
    zipfile="/home/gis/gisdata/data/jupyterlabhub/gitcode/cardanger/zipfile"
    modeltrain="/home/gis/gisdata/data/jupyterlabhub/gitcode/cardanger/modeltrain"
    localtime = time.localtime(time.time())
    shelltext="zip  -r  {}/modeltrain{}_{}_{}.zip  {}/* ".format(zipfile,localtime.tm_mon,localtime.tm_mday,localtime.tm_hour,modeltrain)
    os.system(shelltext)
    
    logdir=os.path.join(modeltrain,"log")
    tfrecorddir=os.path.join(modeltrain,"modelRecord")
    modelcpk=os.path.join(modeltrain,"modelcpk")
    cleardir(logdir)
    cleardir(tfrecorddir)
    cleardir(modelcpk)

    writer=SummaryWriter(tfrecorddir) # 记录训练过程中的情况
    # 加载数据
    csvpath="/home/gis/gisdata/data/jupyterlabhub/gitcode/cardanger/dataset/newcardata/train_test_vail.csv" # 文件路径
    traindataset=MyDataset(csvpath,"train")
    vaildataset=MyDataset(csvpath,"vail")
    testdataset=MyDataset(csvpath,"test")
    trainLoader=DataLoader(dataset=traindataset,batch_size=8,num_workers=8)
    vailLoader=DataLoader(dataset=vaildataset,batch_size=8,num_workers=8)
    testLoader=DataLoader(dataset=testdataset,batch_size=8,num_workers=8)

    # 加载模型--修改模型
    model=models.resnet50(pretrained=True)
    model.fc=nn.Linear(model.fc.in_features,5)
    #model.classifier=nn.Linear(model.classifier.in_features,5)
    model.cuda()

    tlstep,vlstep,ttstep=0,0,0
    minvlac,minttac=0.6,0.6
    traindecary=0.3
    for epoch in range(300):
        model,tlls,tlac,tlstep=train(epoch,tlstep,model,trainLoader,writer,modelcpk,traindecary,lr=0.1)
        model,vlls,vlac,vlstep=test(epoch,vlstep,model,vailLoader,writer,sig="vail")
        model,ttls,ttac,ttstep=test(epoch,ttstep,model,testLoader,writer,sig="test")
        
        # 输出日志
        loginfo(os.path.join(logdir,"train.csv"),tlls,tlac,epoch)
        loginfo(os.path.join(logdir,"vail.csv"),vlls,vlac,epoch)
        loginfo(os.path.join(logdir,"test.csv"),ttls,ttac,epoch)

        # 输出训练信息
        tlac,tlls=printinfo(epoch,tlls,tlac,sig="train")
        vlac,vlls=printinfo(epoch,vlls,vlac,sig="vail")
        ttac,ttls=printinfo(epoch,ttls,ttac,sig="test")
        if tlls<traindecary*3:
            traindecary=traindecary*0.5
        writer.add_scalar('train mean acc',tlac,epoch)
        writer.add_scalar('vail mean acc',vlac,epoch)
        writer.add_scalar('test mean acc',ttac,epoch)
        writer.add_scalar('train mean loss',tlls,epoch)
        writer.add_scalar('vail mean loss',vlls,epoch)
        writer.add_scalar('test mean loss',ttls,epoch)

        if vlac>minvlac:
            minvlac=vlac
            torch.save(model,os.path.join(modelcpk,"model_minvl_{}_{}.pkl".format(int(epoch,minvlac*10000))))
        if ttac>minttac:
            minttac=ttac
            torch.save(model,os.path.join(modelcpk,"model_mintt_{}_{}.pkl".format(int(epoch,minttac*10000))))
        print("current minvailacc:{} mintestacc:{}".format(minvlac,minttac))
    writer.close()
